{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a8bf6335",
   "metadata": {},
   "source": [
    "# ECG CNN PyTorch Demo\n",
    "\n",
    "This notebook demonstrates:\n",
    "\n",
    "0. Imports\n",
    "1. Setup & Config\n",
    "2. Model Definition\n",
    "3. Helper Functions\n",
    "4. Data Loading & Preprocessing\n",
    "5. Grid Search & Training\n",
    "6. Evaluation & Visualization\n",
    "7. External PTB Validation\n",
    "8. Entry Point & Full Runeline: data loading, (subsample), grid search\n",
    "   (targeted), early stopping, evaluation, plotting, PTB validation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5217df5d",
   "metadata": {},
   "source": [
    "## 0. Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6e7edad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import itertools\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from sklearn.metrics import confusion_matrix, classification_report\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Optionally enable AMP for speed on GPU:\n",
    "# from torch.cuda.amp import autocast, GradScaler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9e0f26f",
   "metadata": {},
   "source": [
    "## 1. Setup & Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1fede6e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using 6 CPU threads\n",
      "Running on device: cuda\n"
     ]
    }
   ],
   "source": [
    "# Reproducibility\n",
    "SEED = 22\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "\n",
    "# Thread control\n",
    "torch.set_num_threads(6)\n",
    "print(f\"Using {torch.get_num_threads()} CPU threads\")\n",
    "\n",
    "# Device (GPU if available)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Running on device: {device}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3eca62a6",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 2. Model Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c78ba124",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "class ECGConvNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv1d(1, 64, kernel_size=6)\n",
    "        self.bn1 = nn.BatchNorm1d(64)\n",
    "        self.pool1 = nn.MaxPool1d(3, 2, 1)\n",
    "        self.conv2 = nn.Conv1d(64, 64, 3)\n",
    "        self.bn2 = nn.BatchNorm1d(64)\n",
    "        self.pool2 = nn.MaxPool1d(2, 2, 1)\n",
    "        self.conv3 = nn.Conv1d(64, 64, 3)\n",
    "        self.bn3 = nn.BatchNorm1d(64)\n",
    "        self.pool3 = nn.MaxPool1d(2, 2, 1)\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.fc1 = nn.Linear(64 * 22, 64)\n",
    "        self.fc2 = nn.Linear(64, 32)\n",
    "        self.out = nn.Linear(32, 5)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pool1(F.relu(self.bn1(self.conv1(x))))\n",
    "        x = self.pool2(F.relu(self.bn2(self.conv2(x))))\n",
    "        x = self.pool3(F.relu(self.bn3(self.conv3(x))))\n",
    "        x = self.flatten(x)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        return self.out(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e11f3e",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 3. Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0ed99b62",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def plot_confusion_matrix(cm, classes, normalize=False, title=\"Confusion matrix\"):\n",
    "    if normalize:\n",
    "        sums = cm.sum(axis=1, keepdims=True)\n",
    "        cm = cm.astype(float) / np.where(sums != 0, sums, 1)\n",
    "    plt.imshow(cm, interpolation=\"nearest\", cmap=plt.cm.Blues)\n",
    "    plt.title(title)\n",
    "    plt.colorbar()\n",
    "    ticks = np.arange(len(classes))\n",
    "    plt.xticks(ticks, classes)\n",
    "    plt.yticks(ticks, classes)\n",
    "    thresh = cm.max() / 2.0\n",
    "    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):\n",
    "        val = f\"{cm[i,j]:.2f}\" if normalize else f\"{int(cm[i,j])}\"\n",
    "        plt.text(\n",
    "            j,\n",
    "            i,\n",
    "            val,\n",
    "            ha=\"center\",\n",
    "            va=\"center\",\n",
    "            color=\"white\" if cm[i, j] > thresh else \"black\",\n",
    "        )\n",
    "    plt.ylabel(\"True label\")\n",
    "    plt.xlabel(\"Predicted label\")\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "611ba0f4",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 4. Evaluation & Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7d6fc2ec",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def evaluate_and_plot(\n",
    "    y_true, y_pred, train_accs, val_accs, lr, bs, fold, epochs, out_folder\n",
    "):\n",
    "    print(\n",
    "        f\"\\n=== Final Evaluation (LR={lr}, BS={bs}, Fold={fold}, Epochs={epochs}) ===\"\n",
    "    )\n",
    "    print(\"Classification Report:\")\n",
    "    print(\n",
    "        classification_report(y_true, y_pred, labels=[0, 1, 2, 3, 4], zero_division=0)\n",
    "    )\n",
    "\n",
    "    # Accuracy curve\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.plot(train_accs, label=\"Training\")\n",
    "    ax.plot(val_accs, label=\"Validation\")\n",
    "    ax.set_title(\n",
    "        f\"Model Accuracy by Epoch\\nLR={lr}, BS={bs}, Fold={fold}, Epochs={epochs}\"\n",
    "    )\n",
    "    ax.set_xlabel(\"Epoch\")\n",
    "    ax.set_ylabel(\"Model Accuracy\")\n",
    "    ax.legend()\n",
    "    path = os.path.join(out_folder, f\"final_accuracy_{lr}_{bs}_{fold}_{epochs}.png\")\n",
    "    fig.savefig(path)\n",
    "    plt.close(fig)\n",
    "    print(f\"Saved final accuracy plot to {path}\")\n",
    "\n",
    "    # Confusion matrix\n",
    "    cm = confusion_matrix(y_true, y_pred)\n",
    "    fig, ax = plt.subplots()\n",
    "    plot_confusion_matrix(\n",
    "        cm,\n",
    "        classes=[\"N\", \"S\", \"V\", \"F\", \"Q\"],\n",
    "        normalize=True,\n",
    "        title=f\"Normalized Confusion Matrix\\nLR={lr}, BS={bs}, Fold={fold}, Epochs={epochs}\",\n",
    "    )\n",
    "    path = os.path.join(out_folder, f\"final_confmat_{lr}_{bs}_{fold}_{epochs}.png\")\n",
    "    fig.savefig(path)\n",
    "    plt.close(fig)\n",
    "    print(f\"Saved final confusion matrix to {path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28f8f3d3",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 5. External PTB Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3a7d1c65",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def evaluate_on_ptb(model_path, normal_csv, abnormal_csv, device, out_folder):\n",
    "    print(f\"\\n=== PTB Validation: {model_path} ===\")\n",
    "    ptb_n = pd.read_csv(normal_csv, header=None)\n",
    "    ptb_a = pd.read_csv(abnormal_csv, header=None)\n",
    "    df = pd.concat([ptb_n, ptb_a], ignore_index=True).dropna()\n",
    "    df[187] = df[187].astype(int)\n",
    "    df = df[df[187].isin([0, 1])]\n",
    "    X = df.iloc[:, :-1].values.reshape(-1, 1, 187)\n",
    "    y = df.iloc[:, -1].values\n",
    "    loader = DataLoader(\n",
    "        TensorDataset(\n",
    "            torch.tensor(X, dtype=torch.float32).to(device),\n",
    "            torch.tensor(y, dtype=torch.long),\n",
    "        ),\n",
    "        batch_size=32,\n",
    "    )\n",
    "    ckpt = torch.load(model_path, map_location=device)\n",
    "    model = ECGConvNet().to(device)\n",
    "    model.load_state_dict(ckpt[\"model_state_dict\"])\n",
    "    model.eval()\n",
    "    preds, targets = [], []\n",
    "    with torch.no_grad():\n",
    "        for xb, yb in loader:\n",
    "            out = model(xb)\n",
    "            preds.extend(torch.argmax(out, 1).cpu().tolist())\n",
    "            targets.extend(yb.tolist())\n",
    "    acc = np.mean(np.array(preds) == np.array(targets))\n",
    "    print(f\"PTB Accuracy: {acc:.4f}\")\n",
    "    cm = confusion_matrix(targets, preds)\n",
    "    fig, ax = plt.subplots()\n",
    "    plot_confusion_matrix(\n",
    "        cm[:2, :2],\n",
    "        classes=[\"Normal\", \"Abnormal\"],\n",
    "        normalize=True,\n",
    "        title=\"PTB Confusion Matrix\",\n",
    "    )\n",
    "    path = os.path.join(out_folder, \"ptb_confmat.png\")\n",
    "    fig.savefig(path)\n",
    "    plt.close(fig)\n",
    "    print(f\"Saved PTB confusion matrix to {path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53c12c8f",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 6. Data Loading & Subsample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "061b1a30",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def load_data(data_dir=\"./data\", subsample_frac=0.2):\n",
    "    print(f\"Loading data from {data_dir} and subsampling {subsample_frac*100:.0f}%\")\n",
    "    t = pd.read_csv(os.path.join(data_dir, \"mitbih_train.csv\"), header=None)\n",
    "    s = pd.read_csv(os.path.join(data_dir, \"mitbih_test.csv\"), header=None)\n",
    "    df = pd.concat([t, s], ignore_index=True).dropna()\n",
    "    df[187] = df[187].astype(int)\n",
    "    # subsample\n",
    "    mask = np.random.choice(len(df), size=int(len(df) * subsample_frac), replace=False)\n",
    "    df = df.iloc[mask]\n",
    "    X = df.iloc[:, :-1].values.reshape(-1, 1, df.shape[1] - 1)\n",
    "    y = df.iloc[:, -1].values\n",
    "    print(f\"Final subsample size: {X.shape[0]}\")\n",
    "    return X, y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20860dde",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 7. Grid Search & Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e9c8662b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_grid_search(X, y, out_folder, folds=5):\n",
    "    \"\"\"\n",
    "    Performs grid search with 5-fold CV.\n",
    "    Tracks train/val loss & acc, saves best model, and generates plots.\n",
    "    Returns:\n",
    "      results_df: DataFrame of best val_acc per run\n",
    "      metrics: dict key=(fold,epochs,batch_size,lr) -> metric lists & preds/true\n",
    "    \"\"\"\n",
    "    # param_grid = {\n",
    "    #     \"epochs\": [5, 10],\n",
    "    #     \"batch_size\": [10, 32],\n",
    "    #     \"lr\": [1e-3, 5e-4, 1e-4],\n",
    "    # }\n",
    "    param_grid = {\n",
    "        \"epochs\": [5],\n",
    "        \"batch_size\": [10],\n",
    "        \"lr\": [1e-3, 5e-4],\n",
    "    }\n",
    "\n",
    "    combos = list(\n",
    "        itertools.product(\n",
    "            param_grid[\"epochs\"], param_grid[\"batch_size\"], param_grid[\"lr\"]\n",
    "        )\n",
    "    )\n",
    "    print(f\"Grid search combos: {len(combos)}\")\n",
    "\n",
    "    skf = StratifiedKFold(n_splits=folds, shuffle=True, random_state=SEED)\n",
    "    results = []\n",
    "    metrics = {}\n",
    "\n",
    "    for combo_idx, (epochs, bs, lr) in enumerate(combos, start=1):\n",
    "        print(\n",
    "            f\"=== Combo {combo_idx}/{len(combos)}: epochs={epochs}, bs={bs}, lr={lr} ===\"\n",
    "        )\n",
    "\n",
    "        for fold, (ti, vi) in enumerate(skf.split(X, y), start=1):\n",
    "            print(f\"  Fold {fold}/{folds}\")\n",
    "\n",
    "            X_tr = torch.tensor(X[ti], dtype=torch.float32)\n",
    "            y_tr = torch.tensor(y[ti], dtype=torch.long)\n",
    "            X_vl = torch.tensor(X[vi], dtype=torch.float32)\n",
    "            y_vl = torch.tensor(y[vi], dtype=torch.long)\n",
    "\n",
    "            trn = DataLoader(TensorDataset(X_tr, y_tr), batch_size=bs, shuffle=True)\n",
    "            vl = DataLoader(TensorDataset(X_vl, y_vl), batch_size=bs)\n",
    "\n",
    "            model = ECGConvNet().to(device)\n",
    "            opt = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "            crit = nn.CrossEntropyLoss()\n",
    "\n",
    "            best_acc, best_state = -1.0, None\n",
    "            best_true, best_pred = [], []\n",
    "            train_losses, val_losses = [], []\n",
    "            train_accs, val_accs = [], []\n",
    "\n",
    "            for epoch in range(1, epochs + 1):\n",
    "                # Training\n",
    "                model.train()\n",
    "                running_loss, correct, total = 0.0, 0, 0\n",
    "                for xb, yb in trn:\n",
    "                    xb, yb = xb.to(device), yb.to(device)\n",
    "                    opt.zero_grad()\n",
    "                    out = model(xb)\n",
    "                    loss = crit(out, yb)\n",
    "                    loss.backward()\n",
    "                    opt.step()\n",
    "                    running_loss += loss.item()\n",
    "                    preds = torch.argmax(out, 1)\n",
    "                    correct += (preds == yb).sum().item()\n",
    "                    total += yb.size(0)\n",
    "                avg_train_loss = running_loss / len(trn)\n",
    "                train_acc = correct / total\n",
    "\n",
    "                # Validation\n",
    "                model.eval()\n",
    "                running_vloss, vcorrect, vtotal = 0.0, 0, 0\n",
    "                true, pred = [], []\n",
    "                with torch.no_grad():\n",
    "                    for xb, yb in vl:\n",
    "                        xb, yb = xb.to(device), yb.to(device)\n",
    "                        out = model(xb)\n",
    "                        running_vloss += crit(out, yb).item()\n",
    "                        p = torch.argmax(out, 1)\n",
    "                        vcorrect += (p == yb).sum().item()\n",
    "                        vtotal += yb.size(0)\n",
    "                        true.extend(yb.cpu().tolist())\n",
    "                        pred.extend(p.cpu().tolist())\n",
    "                avg_val_loss = running_vloss / len(vl)\n",
    "                val_acc = vcorrect / vtotal\n",
    "\n",
    "                train_losses.append(avg_train_loss)\n",
    "                val_losses.append(avg_val_loss)\n",
    "                train_accs.append(train_acc)\n",
    "                val_accs.append(val_acc)\n",
    "\n",
    "                print(\n",
    "                    f\"    Epoch {epoch}/{epochs} — loss={avg_train_loss:.4f}/{avg_val_loss:.4f}  \"\n",
    "                    f\"acc={train_acc:.4f}/{val_acc:.4f}\"\n",
    "                )\n",
    "\n",
    "                if val_acc > best_acc:\n",
    "                    best_acc, best_state = val_acc, model.state_dict()\n",
    "                    best_true, best_pred = true.copy(), pred.copy()\n",
    "\n",
    "            # Save best model\n",
    "            fbase = f\"lr{lr:.4f}_bs{bs}_fold{fold}_ep{epochs}\"\n",
    "            model_path = os.path.join(out_folder, f\"model_{fbase}.pt\")\n",
    "            torch.save({\"model_state_dict\": best_state}, model_path)\n",
    "            print(f\"  → Saved model: {model_path}\")\n",
    "\n",
    "            # Plot Loss\n",
    "            fig, ax = plt.subplots()\n",
    "            ax.plot(train_losses, label=\"Training\")\n",
    "            ax.plot(val_losses, label=\"Validation\")\n",
    "            ax.set_title(\n",
    "                f\"Model Loss by Epoch\\nLR={lr}, BS={bs}, Fold={fold}, Epochs={epochs}\"\n",
    "            )\n",
    "            ax.set_xlabel(\"Epoch\")\n",
    "            ax.set_ylabel(\"Model Loss\")\n",
    "            ax.legend()\n",
    "            lp = os.path.join(out_folder, f\"loss_{fbase}.png\")\n",
    "            fig.savefig(lp)\n",
    "            plt.close(fig)\n",
    "            print(f\"  → Saved loss plot: {lp}\")\n",
    "\n",
    "            # Plot Accuracy\n",
    "            fig, ax = plt.subplots()\n",
    "            ax.plot(train_accs, label=\"Training\")\n",
    "            ax.plot(val_accs, label=\"Validation\")\n",
    "            ax.set_title(\n",
    "                f\"Model Accuracy by Epoch\\nLR={lr}, BS={bs}, Fold={fold}, Epochs={epochs}\"\n",
    "            )\n",
    "            ax.set_xlabel(\"Epoch\")\n",
    "            ax.set_ylabel(\"Model Accuracy\")\n",
    "            ax.legend()\n",
    "            ap = os.path.join(out_folder, f\"accuracy_{fbase}.png\")\n",
    "            fig.savefig(ap)\n",
    "            plt.close(fig)\n",
    "            print(f\"  → Saved accuracy plot: {ap}\")\n",
    "\n",
    "            # Plot Confusion Matrix\n",
    "            fig, ax = plt.subplots()\n",
    "            plot_confusion_matrix(\n",
    "                confusion_matrix(best_true, best_pred),\n",
    "                classes=[\"N\", \"S\", \"V\", \"F\", \"Q\"],\n",
    "                normalize=True,\n",
    "                title=(\n",
    "                    f\"Normalized Confusion Matrix\\nLR={lr}, BS={bs}, Fold={fold}, Epochs={epochs}\"\n",
    "                ),\n",
    "            )\n",
    "            cp = os.path.join(out_folder, f\"matrix_{fbase}.png\")\n",
    "            fig.savefig(cp)\n",
    "            plt.close(fig)\n",
    "            print(f\"  → Saved confusion matrix: {cp}\")\n",
    "\n",
    "            # Summary\n",
    "            results.append(\n",
    "                {\n",
    "                    \"fold\": fold,\n",
    "                    \"epochs\": epochs,\n",
    "                    \"batch_size\": bs,\n",
    "                    \"lr\": lr,\n",
    "                    \"val_acc\": best_acc,\n",
    "                    \"model_path\": model_path,\n",
    "                }\n",
    "            )\n",
    "            metrics[(fold, epochs, bs, lr)] = {\n",
    "                \"train_accs\": train_accs,\n",
    "                \"val_accs\": val_accs,\n",
    "                \"train_losses\": train_losses,\n",
    "                \"val_losses\": val_losses,\n",
    "                \"y_true\": best_true,\n",
    "                \"y_pred\": best_pred,\n",
    "            }\n",
    "\n",
    "    return pd.DataFrame(results), metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72249742",
   "metadata": {},
   "source": [
    "## 8. Entry Point & Full Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c881fb07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output folder: ./Outfiles_pytorch\n",
      "Loading data from ./data and subsampling 20%\n",
      "Final subsample size: 21889\n",
      "Grid search combos: 2\n",
      "=== Combo 1/2: epochs=5, bs=10, lr=0.001 ===\n",
      "  Fold 1/5\n",
      "    Epoch 1/5 — loss=0.2029/0.1469  acc=0.9443/0.9587\n",
      "    Epoch 2/5 — loss=0.1255/0.1204  acc=0.9647/0.9667\n",
      "    Epoch 3/5 — loss=0.1026/0.1113  acc=0.9705/0.9724\n",
      "    Epoch 4/5 — loss=0.0895/0.0992  acc=0.9752/0.9744\n",
      "    Epoch 5/5 — loss=0.0791/0.1303  acc=0.9778/0.9687\n",
      "  → Saved model: ./Outfiles_pytorch/model_lr0.0010_bs10_fold1_ep5.pt\n",
      "  → Saved loss plot: ./Outfiles_pytorch/loss_lr0.0010_bs10_fold1_ep5.png\n",
      "  → Saved accuracy plot: ./Outfiles_pytorch/accuracy_lr0.0010_bs10_fold1_ep5.png\n",
      "  → Saved confusion matrix: ./Outfiles_pytorch/matrix_lr0.0010_bs10_fold1_ep5.png\n",
      "  Fold 2/5\n",
      "    Epoch 1/5 — loss=0.2066/0.1425  acc=0.9404/0.9628\n",
      "    Epoch 2/5 — loss=0.1287/0.1293  acc=0.9636/0.9678\n",
      "    Epoch 3/5 — loss=0.1042/0.1045  acc=0.9720/0.9728\n",
      "    Epoch 4/5 — loss=0.0911/0.1003  acc=0.9746/0.9719\n",
      "    Epoch 5/5 — loss=0.0809/0.0953  acc=0.9781/0.9733\n",
      "  → Saved model: ./Outfiles_pytorch/model_lr0.0010_bs10_fold2_ep5.pt\n",
      "  → Saved loss plot: ./Outfiles_pytorch/loss_lr0.0010_bs10_fold2_ep5.png\n",
      "  → Saved accuracy plot: ./Outfiles_pytorch/accuracy_lr0.0010_bs10_fold2_ep5.png\n",
      "  → Saved confusion matrix: ./Outfiles_pytorch/matrix_lr0.0010_bs10_fold2_ep5.png\n",
      "  Fold 3/5\n",
      "    Epoch 1/5 — loss=0.2107/0.1374  acc=0.9415/0.9641\n",
      "    Epoch 2/5 — loss=0.1235/0.1297  acc=0.9656/0.9635\n",
      "    Epoch 3/5 — loss=0.1038/0.1206  acc=0.9716/0.9696\n",
      "    Epoch 4/5 — loss=0.0912/0.0985  acc=0.9744/0.9740\n",
      "    Epoch 5/5 — loss=0.0775/0.1046  acc=0.9767/0.9756\n",
      "  → Saved model: ./Outfiles_pytorch/model_lr0.0010_bs10_fold3_ep5.pt\n",
      "  → Saved loss plot: ./Outfiles_pytorch/loss_lr0.0010_bs10_fold3_ep5.png\n",
      "  → Saved accuracy plot: ./Outfiles_pytorch/accuracy_lr0.0010_bs10_fold3_ep5.png\n",
      "  → Saved confusion matrix: ./Outfiles_pytorch/matrix_lr0.0010_bs10_fold3_ep5.png\n",
      "  Fold 4/5\n",
      "    Epoch 1/5 — loss=0.2150/0.1318  acc=0.9401/0.9635\n",
      "    Epoch 2/5 — loss=0.1286/0.1141  acc=0.9643/0.9669\n",
      "    Epoch 3/5 — loss=0.1062/0.0960  acc=0.9704/0.9714\n",
      "    Epoch 4/5 — loss=0.0892/0.1151  acc=0.9755/0.9701\n",
      "    Epoch 5/5 — loss=0.0784/0.0910  acc=0.9778/0.9751\n",
      "  → Saved model: ./Outfiles_pytorch/model_lr0.0010_bs10_fold4_ep5.pt\n",
      "  → Saved loss plot: ./Outfiles_pytorch/loss_lr0.0010_bs10_fold4_ep5.png\n",
      "  → Saved accuracy plot: ./Outfiles_pytorch/accuracy_lr0.0010_bs10_fold4_ep5.png\n",
      "  → Saved confusion matrix: ./Outfiles_pytorch/matrix_lr0.0010_bs10_fold4_ep5.png\n",
      "  Fold 5/5\n",
      "    Epoch 1/5 — loss=0.2072/0.1566  acc=0.9429/0.9568\n",
      "    Epoch 2/5 — loss=0.1248/0.1282  acc=0.9651/0.9623\n",
      "    Epoch 3/5 — loss=0.1002/0.1101  acc=0.9717/0.9708\n",
      "    Epoch 4/5 — loss=0.0837/0.1063  acc=0.9760/0.9710\n",
      "    Epoch 5/5 — loss=0.0752/0.1145  acc=0.9784/0.9694\n",
      "  → Saved model: ./Outfiles_pytorch/model_lr0.0010_bs10_fold5_ep5.pt\n",
      "  → Saved loss plot: ./Outfiles_pytorch/loss_lr0.0010_bs10_fold5_ep5.png\n",
      "  → Saved accuracy plot: ./Outfiles_pytorch/accuracy_lr0.0010_bs10_fold5_ep5.png\n",
      "  → Saved confusion matrix: ./Outfiles_pytorch/matrix_lr0.0010_bs10_fold5_ep5.png\n",
      "=== Combo 2/2: epochs=5, bs=10, lr=0.0005 ===\n",
      "  Fold 1/5\n",
      "    Epoch 1/5 — loss=0.1969/0.1870  acc=0.9464/0.9529\n",
      "    Epoch 2/5 — loss=0.1224/0.1250  acc=0.9664/0.9678\n",
      "    Epoch 3/5 — loss=0.1051/0.1140  acc=0.9710/0.9708\n",
      "    Epoch 4/5 — loss=0.0919/0.1030  acc=0.9736/0.9730\n",
      "    Epoch 5/5 — loss=0.0771/0.0956  acc=0.9778/0.9740\n",
      "  → Saved model: ./Outfiles_pytorch/model_lr0.0005_bs10_fold1_ep5.pt\n",
      "  → Saved loss plot: ./Outfiles_pytorch/loss_lr0.0005_bs10_fold1_ep5.png\n",
      "  → Saved accuracy plot: ./Outfiles_pytorch/accuracy_lr0.0005_bs10_fold1_ep5.png\n",
      "  → Saved confusion matrix: ./Outfiles_pytorch/matrix_lr0.0005_bs10_fold1_ep5.png\n",
      "  Fold 2/5\n",
      "    Epoch 1/5 — loss=0.2100/0.1177  acc=0.9421/0.9683\n",
      "    Epoch 2/5 — loss=0.1262/0.1096  acc=0.9655/0.9701\n",
      "    Epoch 3/5 — loss=0.1036/0.1153  acc=0.9710/0.9694\n",
      "    Epoch 4/5 — loss=0.0912/0.0849  acc=0.9737/0.9762\n",
      "    Epoch 5/5 — loss=0.0789/0.0939  acc=0.9762/0.9751\n",
      "  → Saved model: ./Outfiles_pytorch/model_lr0.0005_bs10_fold2_ep5.pt\n",
      "  → Saved loss plot: ./Outfiles_pytorch/loss_lr0.0005_bs10_fold2_ep5.png\n",
      "  → Saved accuracy plot: ./Outfiles_pytorch/accuracy_lr0.0005_bs10_fold2_ep5.png\n",
      "  → Saved confusion matrix: ./Outfiles_pytorch/matrix_lr0.0005_bs10_fold2_ep5.png\n",
      "  Fold 3/5\n",
      "    Epoch 1/5 — loss=0.2010/0.1399  acc=0.9431/0.9646\n",
      "    Epoch 2/5 — loss=0.1246/0.1119  acc=0.9669/0.9696\n",
      "    Epoch 3/5 — loss=0.1001/0.1017  acc=0.9697/0.9749\n",
      "    Epoch 4/5 — loss=0.0885/0.1077  acc=0.9739/0.9751\n",
      "    Epoch 5/5 — loss=0.0778/0.1020  acc=0.9777/0.9751\n",
      "  → Saved model: ./Outfiles_pytorch/model_lr0.0005_bs10_fold3_ep5.pt\n",
      "  → Saved loss plot: ./Outfiles_pytorch/loss_lr0.0005_bs10_fold3_ep5.png\n",
      "  → Saved accuracy plot: ./Outfiles_pytorch/accuracy_lr0.0005_bs10_fold3_ep5.png\n",
      "  → Saved confusion matrix: ./Outfiles_pytorch/matrix_lr0.0005_bs10_fold3_ep5.png\n",
      "  Fold 4/5\n",
      "    Epoch 1/5 — loss=0.2045/0.1526  acc=0.9436/0.9598\n",
      "    Epoch 2/5 — loss=0.1233/0.1098  acc=0.9662/0.9705\n",
      "    Epoch 3/5 — loss=0.0994/0.0918  acc=0.9706/0.9760\n",
      "    Epoch 4/5 — loss=0.0833/0.0843  acc=0.9759/0.9762\n",
      "    Epoch 5/5 — loss=0.0751/0.0823  acc=0.9782/0.9785\n",
      "  → Saved model: ./Outfiles_pytorch/model_lr0.0005_bs10_fold4_ep5.pt\n",
      "  → Saved loss plot: ./Outfiles_pytorch/loss_lr0.0005_bs10_fold4_ep5.png\n",
      "  → Saved accuracy plot: ./Outfiles_pytorch/accuracy_lr0.0005_bs10_fold4_ep5.png\n",
      "  → Saved confusion matrix: ./Outfiles_pytorch/matrix_lr0.0005_bs10_fold4_ep5.png\n",
      "  Fold 5/5\n",
      "    Epoch 1/5 — loss=0.1999/0.1493  acc=0.9442/0.9584\n",
      "    Epoch 2/5 — loss=0.1171/0.1155  acc=0.9681/0.9669\n",
      "    Epoch 3/5 — loss=0.0961/0.1084  acc=0.9726/0.9701\n",
      "    Epoch 4/5 — loss=0.0798/0.1053  acc=0.9769/0.9714\n",
      "    Epoch 5/5 — loss=0.0717/0.0957  acc=0.9799/0.9740\n",
      "  → Saved model: ./Outfiles_pytorch/model_lr0.0005_bs10_fold5_ep5.pt\n",
      "  → Saved loss plot: ./Outfiles_pytorch/loss_lr0.0005_bs10_fold5_ep5.png\n",
      "  → Saved accuracy plot: ./Outfiles_pytorch/accuracy_lr0.0005_bs10_fold5_ep5.png\n",
      "  → Saved confusion matrix: ./Outfiles_pytorch/matrix_lr0.0005_bs10_fold5_ep5.png\n",
      "Saved grid search results to results_summary_pytorch.csv/   fold  epochs  batch_size      lr   val_acc  \\\n",
      "0     1       5          10  0.0010  0.974418   \n",
      "1     2       5          10  0.0010  0.973275   \n",
      "2     3       5          10  0.0010  0.975560   \n",
      "3     4       5          10  0.0010  0.975103   \n",
      "4     5       5          10  0.0010  0.970985   \n",
      "5     1       5          10  0.0005  0.973961   \n",
      "6     2       5          10  0.0005  0.976245   \n",
      "7     3       5          10  0.0005  0.975103   \n",
      "8     4       5          10  0.0005  0.978529   \n",
      "9     5       5          10  0.0005  0.973955   \n",
      "\n",
      "                                          model_path  \n",
      "0  ./Outfiles_pytorch/model_lr0.0010_bs10_fold1_e...  \n",
      "1  ./Outfiles_pytorch/model_lr0.0010_bs10_fold2_e...  \n",
      "2  ./Outfiles_pytorch/model_lr0.0010_bs10_fold3_e...  \n",
      "3  ./Outfiles_pytorch/model_lr0.0010_bs10_fold4_e...  \n",
      "4  ./Outfiles_pytorch/model_lr0.0010_bs10_fold5_e...  \n",
      "5  ./Outfiles_pytorch/model_lr0.0005_bs10_fold1_e...  \n",
      "6  ./Outfiles_pytorch/model_lr0.0005_bs10_fold2_e...  \n",
      "7  ./Outfiles_pytorch/model_lr0.0005_bs10_fold3_e...  \n",
      "8  ./Outfiles_pytorch/model_lr0.0005_bs10_fold4_e...  \n",
      "9  ./Outfiles_pytorch/model_lr0.0005_bs10_fold5_e...  \n",
      "Evaluating final model: ./Outfiles_pytorch/model_lr0.0005_bs10_fold4_ep5.pt\n",
      "\n",
      "=== Final Evaluation (LR=0.0005, BS=10, Fold=4, Epochs=5) ===\n",
      "Classification Report:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.98      1.00      0.99      3622\n",
      "           1       1.00      0.72      0.84       110\n",
      "           2       0.94      0.88      0.91       292\n",
      "           3       0.76      0.74      0.75        34\n",
      "           4       0.99      0.98      0.99       320\n",
      "\n",
      "    accuracy                           0.98      4378\n",
      "   macro avg       0.93      0.86      0.89      4378\n",
      "weighted avg       0.98      0.98      0.98      4378\n",
      "\n",
      "Saved final accuracy plot to ./Outfiles_pytorch/final_accuracy_0.0005_10_4_5.png\n",
      "Saved final confusion matrix to ./Outfiles_pytorch/final_confmat_0.0005_10_4_5.png\n",
      "\n",
      "=== PTB Validation: ./Outfiles_pytorch/model_lr0.0005_bs10_fold4_ep5.pt ===\n",
      "PTB Accuracy: 0.2804\n",
      "Saved PTB confusion matrix to ./Outfiles_pytorch/ptb_confmat.png\n",
      "Total runtime: 6.37 minutes\n"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    t0 = time.time()\n",
    "    out_folder = \"./Outfiles_pytorch\"\n",
    "    os.makedirs(out_folder, exist_ok=True)\n",
    "    print(\"Output folder:\", out_folder)\n",
    "\n",
    "    # Load data (with optional subsample)\n",
    "    X, y = load_data(data_dir=\"./data\", subsample_frac=0.2)\n",
    "\n",
    "    # Grid search\n",
    "    results_df, metrics = run_grid_search(X, y, out_folder)\n",
    "    results_path = \"results_summary_pytorch.csv\"\n",
    "    results_df.to_csv(results_path, index=False)\n",
    "    print(f\"Saved grid search results to {results_path}/{results_df}\")\n",
    "\n",
    "    # Final evaluation on best\n",
    "    best_row = results_df.loc[results_df.val_acc.idxmax()]\n",
    "    model_path = best_row.model_path\n",
    "    print(f\"Evaluating final model: {model_path}\")\n",
    "    key = (best_row.fold, best_row.epochs, best_row.batch_size, best_row.lr)\n",
    "    m = metrics[key]\n",
    "\n",
    "    evaluate_and_plot(\n",
    "        y_true=m[\"y_true\"],\n",
    "        y_pred=m[\"y_pred\"],\n",
    "        train_accs=m[\"train_accs\"],\n",
    "        val_accs=m[\"val_accs\"],\n",
    "        lr=best_row.lr,\n",
    "        bs=best_row.batch_size,\n",
    "        fold=best_row.fold,\n",
    "        epochs=best_row.epochs,\n",
    "        out_folder=out_folder,\n",
    "    )\n",
    "\n",
    "    # PTB Evaluation\n",
    "    evaluate_on_ptb(\n",
    "        model_path,\n",
    "        \"./data/ptbdb_normal.csv\",\n",
    "        \"./data/ptbdb_abnormal.csv\",\n",
    "        device,\n",
    "        out_folder,\n",
    "    )\n",
    "\n",
    "    elapsed = (time.time() - t0) / 60\n",
    "    print(f\"Total runtime: {elapsed:.2f} minutes\")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "py:percent,ipynb",
   "main_language": "python",
   "primary": "py:percent"
  },
  "kernelspec": {
   "display_name": "ecg_cnn",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
